from gxl_ai_utils.utils import utils_file

def handle_for_ckwav():
    wav_dir_path = '/home/yanjie/ckwav2'
    txt_dir_path = "/home/yanjie/cktxt"
    wav_path_dict = utils_file.get_scp_for_wav_dir(wav_dir_path,suffix='.wav')
    wav_output_scp = './test_data/ck/wav.scp'
    utils_file.write_dict_to_scp(wav_path_dict,wav_output_scp)
    txt_path_dict = utils_file.get_scp_for_wav_dir(txt_dir_path,suffix='.txt')
    new_txt_dict = {}
    for key,value in txt_path_dict.items():
        new_txt_dict[key] = utils_file.load_first_row_clean(value)
    txt_output_scp = './test_data/ck/text'
    utils_file.write_dict_to_scp(new_txt_dict,txt_output_scp)
    data_list_path = "./test_data/ck/data.list"
    utils_file.do_convert_wav_text_scp_to_jsonl(wav_output_scp,txt_output_scp,data_list_path)


def handle_for_nvwav():
    wav_dir_path = '/home/yanjie/nvwav2'
    txt_dir_path = "/home/yanjie/nvtxt"
    wav_path_dict = utils_file.get_scp_for_wav_dir(wav_dir_path,suffix='.wav')
    wav_output_scp = './test_data/nv/wav.scp'
    utils_file.write_dict_to_scp(wav_path_dict,wav_output_scp)
    txt_path_dict = utils_file.get_scp_for_wav_dir(txt_dir_path,suffix='.txt')
    new_txt_dict = {}
    for key,value in txt_path_dict.items():
        new_txt_dict[key] = utils_file.load_first_row_clean(value)
    txt_output_scp = './test_data/nv/text'
    utils_file.write_dict_to_scp(new_txt_dict,txt_output_scp)
    data_list_path = "./test_data/nv/data.list"
    utils_file.do_convert_wav_text_scp_to_jsonl(wav_output_scp,txt_output_scp,data_list_path)

def handle_for_xlwav():
    wav_dir_path = '/home/yanjie/xlwav2'
    txt_dir_path = "/home/yanjie/xltxt"
    wav_path_dict = utils_file.get_scp_for_wav_dir(wav_dir_path,suffix='.wav')
    wav_output_scp = './test_data/xl/wav.scp'
    utils_file.write_dict_to_scp(wav_path_dict,wav_output_scp)
    txt_path_dict = utils_file.get_scp_for_wav_dir(txt_dir_path,suffix='.txt')
    new_txt_dict = {}
    for key,value in txt_path_dict.items():
        new_txt_dict[key] = utils_file.load_first_row_clean(value)
    txt_output_scp = './test_data/xl/text'
    utils_file.write_dict_to_scp(new_txt_dict,txt_output_scp)
    data_list_path = "./test_data/xl/data.list"
    utils_file.do_convert_wav_text_scp_to_jsonl(wav_output_scp,txt_output_scp,data_list_path)


def handle_for_metric():
    wav_dir_path = '/home/yanjie/asr_metric_test_data'
    txt_path = "/home/yanjie/asr_metric_wav2txt.txt"
    wav_path_dict = utils_file.get_scp_for_wav_dir(wav_dir_path,suffix='.wav', recursive=True)
    wav_output_scp = './test_data/asr_metric_test_data/wav.scp'
    txt_output_scp = './test_data/asr_metric_test_data/text'
    utils_file.copy_file(txt_path,txt_output_scp)
    text_dict = utils_file.load_dict_from_scp(txt_output_scp)
    new_text_dict = {}
    new_wav_dict = {}
    for key , value in text_dict.items():
        new_key = key.replace('.wav','')
        new_text_dict[new_key] = value
        for key_wav in wav_path_dict.keys():
            if key_wav.startswith(new_key):
                new_wav_dict[new_key] = wav_path_dict[key_wav]
                break
    utils_file.write_dict_to_scp(new_wav_dict,wav_output_scp)
    utils_file.write_dict_to_scp(new_text_dict,txt_output_scp)

    data_list_path = "./test_data/asr_metric_test_data/data.list"
    utils_file.do_convert_wav_text_scp_to_jsonl(wav_output_scp,txt_output_scp,data_list_path)
if __name__ == '__main__':
    handle_for_metric()
